{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating your own pdf\n",
    "\n",
    "A core feature of zfit is the ability to create custom pdfs and functions in an simple and straightforward way.\n",
    "\n",
    "There are two main possibilities to create a custom pdf, an easier for most use-cases and an advanced way..\n",
    "\n",
    "## The simple way\n",
    "\n",
    "While the same works for functions, an example with a PDF is shown here.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import zfit\n",
    "from zfit import z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first way is the most simple and should only be used for the trivial cases, i.e. if you're not familiar with Python classes (especially not with the `__init__` method)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyGauss(zfit.pdf.ZPDF):\n",
    "    _N_OBS = 1  # dimension, can be omitted\n",
    "    _PARAMS = ['mean', 'std']  # the name of the parameters\n",
    "\n",
    "    @zfit.supports()\n",
    "    def _unnormalized_pdf(self, x, params):\n",
    "        x0 = x[0]  # using the 0th axis\n",
    "        mean = params['mean']\n",
    "        std = params['std']\n",
    "        return z.exp(- ((x0 - mean) / std) ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Done. Now we can use our pdf already!\n",
    "\n",
    "The slightly more general way involves overwritting the `__init__` and gives you all the possible flexibility: to use custom parameters, to preprocess them etc.\n",
    "\n",
    "Here we inherit from `BasePDF`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyGauss(zfit.pdf.BasePDF):\n",
    "\n",
    "    def __init__(self, mean, std, obs, extended=None, norm=None, name=None, label=None):\n",
    "        params = {'mean': mean,  # 'mean' is the name as it will be named in the PDF, mean is just the parameter to create the PDF\n",
    "                  'std': std\n",
    "                  }\n",
    "        super().__init__(obs=obs, params=params, extended=extended, norm=norm,\n",
    "                         name=name, label=label)\n",
    "\n",
    "    @zfit.supports()\n",
    "    def _unnormalized_pdf(self, x, params):\n",
    "        x0 = x[0]  # using the 0th axis\n",
    "        mean = params['mean']\n",
    "        std = params['std']\n",
    "        return z.exp(- ((x0 - mean) / std) ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = zfit.Space('obs1', -3, 6)\n",
    "\n",
    "data_np = np.random.random(size=1000)\n",
    "data = zfit.Data(data_np, obs=obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create two parameters and an instance of your own pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = zfit.Parameter(\"mean\", 1.)\n",
    "std = zfit.Parameter(\"std\", 1.)\n",
    "my_gauss = MyGauss(obs=obs, mean=mean, std=std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probs = my_gauss.pdf(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(probs[:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to make sure it's a numpy array, we can use `zfit.run`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could improve our PDF by registering an integral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gauss_integral_from_any_to_any(limits, params, model):\n",
    "    lower, upper = limits.v1.limits\n",
    "    mean = params['mean']\n",
    "    std = params['std']\n",
    "    # write your integral here\n",
    "    return 42.  # dummy integral, must be a scalar!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "limits = zfit.Space(axes=0, lower=zfit.Space.ANY_LOWER, upper=zfit.Space.ANY_UPPER)\n",
    "MyGauss.register_analytic_integral(func=gauss_integral_from_any_to_any, limits=limits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More advanced custom PDFs are introduced in the guide on [custom PDFs](custom_pdfs.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}